# -*- coding: utf-8 -*-
"""
中性策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
import os
import ccxt
import pandas as pd
import numpy as np
from glob import glob
from datetime import datetime
from api.market import load_market, fetch_binance_candle_data, fetch_binance_funding_rate
from utils.commons import retry_wrapper
from config import data_center_path, exchange_basic_config
from tqdm import tqdm

# 获取脚本文件的路径
script_path = os.path.abspath(__file__)

# 提取文件名
script_filename = os.path.basename(script_path).split('.')[0]


def download(run_time):
    """
    根据获取数据的情况，自行编写下载数据函数
    :param run_time:    运行时间
    """
    print(f'执行{script_filename}脚本 download 开始')
    _time = datetime.now()
    # =====创建公共的exchange
    exchange = ccxt.binance(exchange_basic_config)

    # =====获取所有币种的K线
    interval = '1h'  # 根据持仓周期自动调整获取1小时还是1天的k线
    # ===获取合约k线数据
    # =获取U本位合约交易对的信息
    swap_symbol_list = load_market(exchange, symbol_type='swap')[0]

    s_time = datetime.now()
    # 获取合约k线数据 & 保存
    _fetch_candle_data_and_save(exchange, swap_symbol_list, run_time, interval, 'swap')
    print('获取所有合约币种K线数据 & 保存完成，花费时间：', datetime.now() - s_time)

    # ===获取现货k线数据
    # =加载现货交易对信息
    spot_symbol_list = load_market(exchange, symbol_type='spot')[0]

    s_time = datetime.now()
    # 获取现货k线数据 & 保存
    _fetch_candle_data_and_save(exchange, spot_symbol_list, run_time, interval, 'spot')
    print('获取所有现货币种K线数据 & 保存完成，花费时间：', datetime.now() - s_time)

    print(f'执行{script_filename}脚本 download 完成', datetime.now() - _time)


def clean_data():
    """
    根据获取数据的情况，自行编写清理冗余数据函数
    """
    print(f'执行{script_filename}脚本 clear_duplicates 开始')
    _time = datetime.now()
    # 遍历合约和现货目录
    for symbol_type in ['swap', 'spot']:
        # 获取目录路径
        save_path = os.path.join(data_center_path, script_filename, symbol_type)
        # 获取.csv结尾的文件目录
        file_list = glob(os.path.join(save_path, '*.csv'))
        # 遍历文件进行操作
        for _file in file_list:
            # 文件存在，去重之后重新保存
            if os.path.exists(_file):
                df = pd.read_csv(_file, encoding='gbk', parse_dates=['candle_begin_time'])  # 读取本地数据
                df.drop_duplicates(subset=['candle_begin_time'], keep='last', inplace=True)  # 去重保留最新的数据
                df.sort_values('candle_begin_time', inplace=True)  # 通过candle_begin_time排序
                df = df[-2400:]  # 保留最近2400根k线，防止数据堆积过多(2400根，大概100天数据)
                df.to_csv(_file, encoding='gbk', index=False)  # 保存文件

    print(f'执行{script_filename}脚本 clear_duplicates 完成', datetime.now() - _time)


def _fetch_candle_data_and_save(exchange, symbol_list, run_time, interval, symbol_type):
    """
    下载脚本内部使用的函数，不对外使用
    获取k线数据，并保存到本地
    :param exchange: 交易所对象
    :param symbol_list: 需要获取数据的币种列表
    :param run_time: 运行时间
    :param interval: 获取数据周期
    :param symbol_type: 币种类型
    """
    # ===在data目录下创建当前脚本存放数据的目录
    save_path = os.path.join(data_center_path, script_filename, symbol_type)
    # 目录不存在就创建
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # ===获取最新资金费数据
    if symbol_type == 'swap':
        last_funding_df = retry_wrapper(exchange.fapiPublic_get_premiumindex, func_name='获取最新的资金费数据')
        last_funding_df = pd.DataFrame(last_funding_df)
        last_funding_df['nextFundingTime'] = pd.to_datetime(last_funding_df['nextFundingTime'], unit='ms')
        last_funding_df['time'] = pd.to_datetime(last_funding_df['time'], unit='ms')
        last_funding_df = last_funding_df[['symbol', 'nextFundingTime', 'lastFundingRate']]  # 保留部分字段
        last_funding_df.rename(columns={'nextFundingTime': 'fundingTime', 'lastFundingRate': 'fundingRate'}, inplace=True)

    # ===遍历获取币种数据 & 保存
    for symbol in tqdm(symbol_list):
        # =构建存储文件的路径
        _file_path = os.path.join(save_path, f'{symbol}.csv')
        # =判断文件是否存在
        if_file_exists = os.path.exists(_file_path)
        # =获取合约k线请求的数量
        # 如果存在目录，表示已经有文件存储，默认获取50根k线(为什么是50?拍脑袋决定的。这里你可以调整，不建议超过99，会增加请求权重)
        # 如果不存在目录，表示首次运行，获取1500根k线
        limit = 50 if if_file_exists else 1500

        # =获取k线数据
        df = fetch_binance_candle_data(exchange, symbol, run_time, limit, interval=interval, symbol_type=symbol_type)
        # 返回None或者空的df，不放到result里
        if df is None or df.empty:
            continue

        # =追加部分数据
        if symbol_type == 'spot':
            df['fundingRate'] = np.nan
        else:
            # =获取资金费数据请求的数量
            # 如果存在目录，表示已经有文件存储，默认获取45条资金费数据(为什么是45条？拍脑袋的，45条数据就是15天)
            # 如果不存在目录，表示首次运行，获取1000条资金费数据
            limit = 45 if if_file_exists else 1000
            # =获取历史资金费数据
            fund_df = fetch_binance_funding_rate(exchange, symbol, run_time, limit)
            # 返回None或者空的df，跳过
            if fund_df is None or fund_df.empty:
                df['fundingRate'] = np.nan
            else:
                # 合并最新数据
                symbol_df = last_funding_df[last_funding_df['symbol'] == symbol]  # 筛选当前币种的最新资金费数据
                fund_df = pd.concat([fund_df, symbol_df], ignore_index=True)  # 数据合并
                fund_df.drop_duplicates(subset=['fundingTime'], keep='last', inplace=True)  # 去重保留最新的数据

                # 将数据合并到k线上
                df = pd.merge(df, fund_df[['fundingTime', 'fundingRate']], left_on=['candle_begin_time'], right_on=['fundingTime'], how='left')
                del df['fundingTime']

        # =保存数据
        # 路径存在，数据直接追加
        if if_file_exists:
            df = df[-10:]  # 指定最新的10条数据去追加(如果你的数据停了很久怎么办？直接删了data_center重跑吧)
            df.to_csv(_file_path, encoding='gbk', index=False, header=False, mode='a')
        else:
            df.to_csv(_file_path, encoding='gbk', index=False)
